"""
Project: RadarBook
File: linear_array.py
Created by: Lee A. Harrison
On: 7/31/2018
Created with: PyCharm

Copyright (C) 2019 Artech House (artech@artechhouse.com)
This file is part of Introduction to Radar Using Python and MATLAB
and can not be copied and/or distributed without the express permission of Artech House.
"""
import warnings
from typing import Tuple
import numpy as np
from scipy.constants import c, pi
from numpy import cos, floor, roll, amax, ones
from scipy.special import binom
from scipy.signal.windows import chebwin, hann, hamming, blackmanharris, kaiser
import torch
import torch.nn.functional as F


def array_factor(number_of_elements, scan_angle, element_spacing, frequency, theta, window_type, side_lobe_level):
    """
    Calculate the array factor for a linear binomial excited array.
    :param window_type: The string name of the window.
    :param side_lobe_level: The sidelobe level for Tschebyscheff window (dB).
    :param number_of_elements: The number of elements in the array.
    :param scan_angle: The angle to which the main beam is scanned (rad).
    :param element_spacing: The distance between elements.
    :param frequency: The operating frequency (Hz).
    :param theta: The angle at which to evaluate the array factor (rad).
    :return: The array factor as a function of angle.
    """
    # Calculate the wavenumber
    k = 2.0 * pi * frequency / c

    # Calculate the phase
    psi = k * element_spacing * (cos(theta) - cos(scan_angle))
    print(f'psi: {psi.shape};')

    # Calculate the coefficients
    if window_type == 'Uniform':
        coefficients = ones(number_of_elements)
    elif window_type == 'Binomial':
        coefficients = binom(number_of_elements-1, range(0, number_of_elements))
    elif window_type == 'Tschebyscheff':
        warnings.simplefilter("ignore", UserWarning)
        coefficients = chebwin(number_of_elements, at=side_lobe_level, sym=True)
    elif window_type == 'Kaiser':
        coefficients = kaiser(number_of_elements, 6, True)
    elif window_type == 'Blackman-Harris':
        coefficients = blackmanharris(number_of_elements, True)
    elif window_type == 'Hanning':
        coefficients = hann(number_of_elements, True)
    elif window_type == 'Hamming':
        coefficients = hamming(number_of_elements, True)

    print(f'??????????????? coefficients: {coefficients.shape}; \n{coefficients};')

    # Calculate the offset for even/odd 21 =》10
    offset = int(floor(number_of_elements / 2))
    # coefficients: (N,)
    # Odd case
    if number_of_elements & 1:
        coefficients = roll(coefficients, offset + 1)
        coefficients[0] *= 0.5
        # 因为\theta划分了1000份，所以af有1000个元素：[11, 1000] => [1000]
        af = sum(coefficients[i] * cos(i * psi) for i in range(offset + 1))
        return af / amax(abs(af))
    # Even case
    else:
        coefficients = roll(coefficients, offset) # 把后一半倒到前面
        # coefficients[i] * cos((i + 0.5) * psi) for i in range(offset): (1000, 16)
        af = sum(coefficients[i] * cos((i + 0.5) * psi) for i in range(offset)) # af: (1000,)
        return af / amax(abs(af)) # 对其进行归一化

def calculate_bw_sll(af:np.ndarray, x_scale:float) -> Tuple[float, float]:
    # 找到主瓣的最大值及其索引
    mainlobe_max_index = np.argmax(af)
    mainlobe_max_value = af[mainlobe_max_index]
    # 找到-3dB点
    bw_level = -3.0 # 找-3dB的索引值
    left_idx = np.argmin(np.abs(af[0:mainlobe_max_index] - bw_level))
    right_idx = np.argmin(np.abs(af[mainlobe_max_index:] - bw_level)) + mainlobe_max_index
    beamwidth = (right_idx - left_idx) * x_scale * 180 / np.pi # 以角度为单位
    # 
    sidelobe_level_db = mainlobe_max_value
    dec_num = 0
    # 下降阶段
    dec_idx = 0
    for idx in range(mainlobe_max_index, 0, -1):
        if af[idx-1] > af[idx]:
            dec_idx = idx
            break
    sll_idx = 0
    for idx in range(dec_idx, 0, -1):
        if af[idx-1] < af[idx]:
            sll_idx = idx
            break
    sidelobe_level_db = af[sll_idx]
    print(f'旁瓣：{sll_idx}: {sidelobe_level_db}; ????????')
    # 计算波束宽度方法2
    aft = torch.from_numpy(af*-1)
    aft = torch.max(aft) - aft
    aft = F.relu(aft - (torch.max(aft) - 3))
    idxs = torch.nonzero(aft)
    left_idx = torch.min(idxs)
    right_idx = torch.max(idxs)
    bw = (right_idx - left_idx) * x_scale * 180 / np.pi # 以角度为单位
    v1 = -F.softmax(aft, dim=0)
    print(f'????? v1: {af.shape}; vs {np.sum(af)};')
    return beamwidth, sidelobe_level_db

def dl_array_factor(number_of_elements, scan_angle, element_spacing, frequency, theta, coefficients):
    """
    Calculate the array factor for a linear binomial excited array.
    :param window_type: The string name of the window.
    :param side_lobe_level: The sidelobe level for Tschebyscheff window (dB).
    :param number_of_elements: The number of elements in the array.
    :param scan_angle: The angle to which the main beam is scanned (rad).
    :param element_spacing: The distance between elements.
    :param frequency: The operating frequency (Hz).
    :param theta: The angle at which to evaluate the array factor (rad).
    :return: The array factor as a function of angle.
    """
    # Calculate the wavenumber
    c = 3E8
    k = 2.0 * torch.pi * frequency / c

    # Calculate the phase
    psi = k * element_spacing * (torch.cos(theta) - cos(scan_angle))
    # Calculate the offset for even/odd 21 =》10
    offset = int(floor(number_of_elements / 2))
    # coefficients: (N,)
    # Odd case
    if number_of_elements & 1:
        coefficients = torch.roll(coefficients, offset + 1)
        coefficients[0] *= 0.5
        # 因为\theta划分了1000份，所以af有1000个元素：[11, 1000] => [1000]
        af = None
        for i in range(offset):
            if af is None:
                af = coefficients[i] * torch.cos((i + 0.5) * psi)
            else:
                af += coefficients[i] * torch.cos((i + 0.5) * psi)
        # af = torch.sum(coefficients[i] * cos(i * psi) for i in range(offset + 1))
        return af / torch.amax(torch.abs(af))
    # Even case
    else:
        coefficients = torch.roll(coefficients, offset) # 把后一半倒到前面
        af = None
        for i in range(offset):
            if af is None:
                af = coefficients[i] * torch.cos((i + 0.5) * psi)
            else:
                af += coefficients[i] * torch.cos((i + 0.5) * psi)
        # af = torch.sum(coefficients[i] * cos((i + 0.5) * psi) for i in range(offset)) # af: (1000,)
        return af / torch.amax(torch.abs(af)) # 对其进行归一化    
